# import cv2
# import numpy as np
# import os
# import glob


# def calculate_centroids(mask):
#     """计算标注图中各目标的质心坐标，背景为0，目标ID为1,2,3..."""
#     results = []
#     breakpoint()
#     # 获取所有唯一的目标ID（排除背景0）
#     unique_labels = np.unique(mask)
#     unique_labels = unique_labels[unique_labels > 0]  # 排除背景0
    
#     # 遍历所有目标ID
#     for label in unique_labels:
#         # 找到当前目标的所有像素位置
#         y, x = np.where(mask == label) # x y 是数组
#         if len(y) > 0:  # 确保目标存在
#             # 计算质心坐标
#             centroid_row = np.mean(y).round(2)
#             centroid_col = np.mean(x).round(2)
#             results.append((label, centroid_row, centroid_col))
    
#     return results


# def process_sequence(seq_folder, output_txt):
#     # 在序列文件夹下添加mask子目录
#     mask_folder = os.path.join(seq_folder, 'mask')
#     mask_files = sorted(glob.glob(os.path.join(mask_folder, '*.png')))
#     if not mask_files:
#         print(f"警告：在 {mask_folder} 中未找到PNG文件")
#         return

#     with open(output_txt, 'w') as f:
#         for frame_idx, mask_path in enumerate(mask_files, start=1):
#             mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
#             if mask is None:
#                 print(f"警告：无法读取图像 {mask_path}")
#                 continue

#             # 直接处理标注图，不需要二值化
#             centroids = calculate_centroids(mask)
#             num_targets = len(centroids)

#             line = f"{frame_idx:05d}   {num_targets}   "
#             # 按ID排序保证输出顺序一致（可选）
#             centroids.sort(key=lambda x: x[0])
#             for label, row, col in centroids:
#                 line += f"{label}     {row:.2f}  {col:.2f}   "

#             f.write(line.strip() + '\n')
#             print(f"已处理：{os.path.basename(mask_path)} - 检测到 {num_targets} 个目标")




import cv2  # 添加缺失的导入
import numpy as np
import os
import glob

def calculate_centroids(mask):
    """计算标注图中各目标的质心坐标，背景为0，目标ID为1,2,3..."""
    results = []
    # 移除 breakpoint() 调试语句
    
    # 获取所有唯一的目标ID（排除背景0）
    unique_labels = np.unique(mask)
    unique_labels = unique_labels[unique_labels > 0]  # 排除背景0
    
    # 遍历所有目标ID
    for label in unique_labels:
        # 找到当前目标的所有像素位置
        y, x = np.where(mask == label)
        if len(y) > 0:  # 确保目标存在
            # 计算质心坐标
            centroid_row = np.mean(y).round(2)
            centroid_col = np.mean(x).round(2)
            results.append((label, centroid_row, centroid_col))
    
    return results

def process_sequence(seq_folder, output_txt):
    # 在序列文件夹下添加mask子目录
    mask_folder = os.path.join(seq_folder, 'mask')
    mask_files = sorted(glob.glob(os.path.join(mask_folder, '*.png')))
    if not mask_files:
        print(f"警告：在 {mask_folder} 中未找到PNG文件")
        return

    # 添加目标跟踪字典来维持轨迹连续性
    track_id_counter = 1
    target_tracks = {}  # {original_id: track_id}
    
    with open(output_txt, 'w') as f:
        for frame_idx, mask_path in enumerate(mask_files, start=1):
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            if mask is None:
                print(f"警告：无法读取图像 {mask_path}")
                continue

            # 直接处理标注图，不需要二值化
            centroids = calculate_centroids(mask)
            
            # 为新目标分配连续的轨迹ID
            tracked_centroids = []
            for label, row, col in centroids:
                if label not in target_tracks:
                    target_tracks[label] = track_id_counter
                    track_id_counter += 1
                track_id = target_tracks[label]
                tracked_centroids.append((track_id, row, col))
            
            num_targets = len(tracked_centroids)
            line = f"{frame_idx:05d}   {num_targets}   "
            
            # 按轨迹ID排序保证输出顺序一致
            tracked_centroids.sort(key=lambda x: x[0])
            for track_id, row, col in tracked_centroids:
                line += f"{track_id}     {row:.2f}  {col:.2f}   "

            f.write(line.strip() + '\n')
            print(f"已处理：{os.path.basename(mask_path)} - 检测到 {num_targets} 个目标")

# 主处理流程
if __name__ == '__main__':
    root_dir = r"./datasets/SatVideoIRSDT/val"
    output_base_dir = r"./evaluate/true_val_centroids"
    os.makedirs(output_base_dir, exist_ok=True)
    # 不再需要阈值参数，因为直接处理标注图

    seq_folders = sorted(glob.glob(os.path.join(root_dir, '0*')))
    if not seq_folders:
        print(f"错误：在 {root_dir} 中未找到序列文件夹")
    else:
        print(f"找到 {len(seq_folders)} 个序列文件夹，开始处理...")
        for seq_folder_path in seq_folders:
            seq_name = os.path.basename(seq_folder_path)
            output_txt_path = os.path.join(output_base_dir, f"{seq_name}.txt")

            if os.path.exists(seq_folder_path):
                print(f"\n处理序列: {seq_name}")
                os.makedirs(os.path.dirname(output_txt_path), exist_ok=True)
                process_sequence(seq_folder_path, output_txt_path)
                print(f"处理完成！结果已保存至: {output_txt_path}")
            else:
                print(f"警告：序列 {seq_name} 文件夹不存在")
        print("\n所有序列处理完成！")